import dask.dataframe as dd
import pandas as pd
from datetime import datetime
import sqlite3
import psycopg2
from psycopg2 import sql
import pymysql
import os
import traceback
from sqlalchemy import create_engine
from sqlalchemy.engine import URL
from urllib.parse import quote_plus
import sys

# 增加系统递归深度限制（临时解决方案）
sys.setrecursionlimit(10000)

def get_db_connection(database_config):
    """根据配置获取数据库连接"""
    db_type = database_config.get('type', 'sqlite')

    if db_type == 'sqlite':
        return sqlite3.connect(database_config['database'])

    elif db_type == 'mysql':
        return pymysql.connect(
            host=database_config['host'],
            user=database_config['user'],
            password=database_config['password'],
            database=database_config['database'],
            port=database_config.get('port', 3306),
            charset=database_config.get('charset', 'utf8mb4')
        )

    elif db_type == 'postgresql':
        print(f"尝试连接PostgreSQL数据库: {database_config['host']}:{database_config['port']}/{database_config['database']}")
        return psycopg2.connect(
            host=database_config['host'],
            user=database_config['user'],
            password=database_config['password'],
            database=database_config['database'],
            port=database_config.get('port', 5432)
        )

    else:
        raise ValueError(f"不支持的数据库类型: {db_type}")


def execute_query(connection, query, params=None):
    """执行SQL查询并返回结果"""
    try:
        with connection.cursor() as cursor:
            cursor.execute(query, params)
            if query.strip().lower().startswith('select'):
                columns = [desc[0] for desc in cursor.description]
                data = cursor.fetchall()
                return pd.DataFrame(data, columns=columns)
            else:
                connection.commit()
                return cursor.rowcount
    except Exception as e:
        print(f"执行SQL查询时出错: {e}")
        traceback.print_exc()  # 打印完整堆栈信息
        raise


def vectorized_calculate_age(birth_date_str_series):
    """
    向量化年龄计算函数，处理多种日期格式
    """
    # 转换为字符串并去除空格
    birth_date_str_series = birth_date_str_series.fillna('').astype(str).str.strip()
    
    # 定义常见日期格式列表
    date_formats = [
        '%Y-%m-%d %H:%M:%S',
        '%Y-%m-%d',
        '%m/%d/%Y',
        '%d-%m-%Y',
        '%Y%m%d',
        '%Y年%m月%d日',
        '%Y-%m',
        '%Y%m',
        '%Y'
    ]
    
    # 初始化结果列表
    ages = []
    
    for date_str in birth_date_str_series:
        if not date_str:
            ages.append(None)
            continue
            
        birth_date = None
        for fmt in date_formats:
            try:
                birth_date = datetime.strptime(date_str, fmt)
                break
            except ValueError:
                continue
                
        if birth_date:
            today = datetime(2025, 1, 1)
            age = today.year - birth_date.year - ((today.month, today.day) < (birth_date.month, birth_date.day))
            
            # 检查年龄是否在合理范围内
            if age < 0 or age > 100:
                print(f"警告: 计算出的年龄超出合理范围: {age} (原始值: {date_str})")
                ages.append(None)
            else:
                ages.append(age)
        else:
            print(f"警告: 无法解析的出生日期格式: {date_str}")
            ages.append(None)
    
    return pd.Series(ages)


def categorize_age(age):
    """将年龄划分为指定的年龄段"""
    if pd.isna(age):
        return '未知年龄'
    try:
        age = int(age)
        if age <= 6:
            return '≤6'
        elif 6 < age <= 10:
            return '6＜年龄≤10'
        elif 10 < age <= 18:
            return '10＜年龄≤18'
        elif 18 < age <= 30:
            return '18＜年龄≤30'
        elif 30 < age <= 40:
            return '30＜年龄≤40'
        elif 40 < age <= 50:
            return '40＜年龄≤50'
        elif 50 < age <= 60:
            return '50＜年龄≤60'
        elif 60 < age <= 65:
            return '60＜年龄≤65'
        elif 65 < age <= 70:
            return '65＜年龄≤70'
        elif 70 < age <= 80:
            return '70＜年龄≤80'
        elif 80 < age <= 90:
            return '80＜年龄≤90'
        elif 90 < age <= 100:
            return '90＜年龄≤100'
        else:
            return '100＜年龄'
    except Exception as e:
        print(f"分类年龄时出错: {age}, 错误: {e}")
        traceback.print_exc()
        return '未知年龄'


def extract_year(date_str):
    """增强版年份提取函数，处理多种日期时间格式"""
    if pd.isna(date_str):
        print(f"警告: 日期字段为空: {date_str}")
        return None

    date_str = str(date_str).strip()

    # 检查日期字符串长度是否符合基本要求
    if len(date_str) < 4:
        print(f"警告: 日期字段长度不足4位: {date_str}")
        return None

    # 尝试解析日期时间格式
    try:
        # 尝试常见的日期时间格式
        for fmt in ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d', '%Y%m%d', '%m/%d/%Y'):
            try:
                dt = datetime.strptime(date_str, fmt)
                return dt.year
            except ValueError:
                continue

        # 如果所有格式都不匹配，尝试直接截取前4位
        year_str = date_str[:4]
        if year_str.isdigit():
            year = int(year_str)
            # 检查年份是否在合理范围内
            if 1900 <= year <= 2099:
                return year
            else:
                print(f"警告: 年份超出合理范围: {year} (原始值: {date_str})")
                return None

        print(f"警告: 无法解析的日期格式: {date_str}")
        return None

    except Exception as e:
        print(f"错误: 处理日期字段时发生异常: {date_str}, 异常: {e}")
        traceback.print_exc()
        return None


def analyze_data(database_config, output_file=None, chunksize=100000):
    """从数据库读取大数据并进行分析（仅使用kdp01表数据）"""
    try:
        print(f"开始分析数据，数据库类型: {database_config['type']}")
        
        # 获取数据库连接 - 使用SQLAlchemy URL对象构建连接，避免密码特殊字符问题
        db_type = database_config['type']
        
        if db_type == 'sqlite':
            engine = create_engine(f"sqlite:///{database_config['database']}")
        elif db_type == 'mysql':
            # 使用URL对象构建MySQL连接
            url = URL.create(
                drivername="mysql+pymysql",
                username=database_config['user'],
                password=database_config['password'],
                host=database_config['host'],
                port=database_config.get('port', 3306),
                database=database_config['database']
            )
            engine = create_engine(url)
        elif db_type == 'postgresql':
            # 使用URL对象构建PostgreSQL连接
            url = URL.create(
                drivername="postgresql",
                username=database_config['user'],
                password=database_config['password'],
                host=database_config['host'],
                port=database_config.get('port', 5432),
                database=database_config['database']
            )
            engine = create_engine(url)
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")

        # 创建数据库连接对象
        connection = get_db_connection(database_config)

        # 执行基本统计查询
        queries = {
            '注册人数': "SELECT COUNT(1) as count FROM kdp01 WHERE axcp0077='0'",
            '身份证重复人数': """
                SELECT SUM(num) as count FROM (
                    SELECT COUNT(1) as num 
                    FROM kdp01 k 
                    WHERE axcp0004='01' AND axcp0005 != '' 
                    GROUP BY axcp0005 
                    HAVING COUNT(1) > 1
                ) t
            """,
            '身份证缺失人数': "SELECT COUNT(1) as count FROM kdp01 k WHERE axcp0004='01' AND (axcp0005 IS NULL OR axcp0005 = '')",
            '有效身份证件用户数': "SELECT COUNT(1) as count FROM kdp01 WHERE axcp0004='01' AND LENGTH(axcp0005) = 18",
            '非身份证件用户数': "SELECT COUNT(1) as count FROM kdp01 WHERE axcp0004!='01'",
            '志愿时长为0的用户数': "SELECT COUNT(1) as count FROM kdp01 WHERE albp0289='0.0' OR albp0289 IS NULL",
        }

        basic_stats = {}
        for name, query in queries.items():
            try:
                print(f"执行统计查询: {name}")
                result = execute_query(connection, query)
                basic_stats[name] = result.iloc[0]['count']
            except Exception as e:
                print(f"执行统计查询 '{name}' 时出错: {e}")
                traceback.print_exc()
                basic_stats[name] = None

        # 读取kdp01表数据
        try:
            print(f"开始读取kdp01表数据，分块大小: {chunksize}")
            columns_str = ', '.join([
                'axcp0039', 'axcp0004', 'axcp0005', 'albp0289', 'axcp0077'
            ])
            query = f"SELECT {columns_str} FROM kdp01"

            dfs = []
            for i, chunk in enumerate(pd.read_sql(query, engine, chunksize=chunksize)):
                print(f"已读取数据块 {i+1}")
                dfs.append(dd.from_pandas(chunk, npartitions=1))

            if not dfs:
                raise ValueError("无法从kdp01表中读取任何数据")

            detail_data = dd.concat(dfs)
            print(f"成功读取{len(detail_data)}条记录")
        except Exception as e:
            print(f"读取kdp01表时出错: {e}")
            traceback.print_exc()
            raise

        # 修复列名
        try:
            detail_data = detail_data.rename(columns={
                'axcp0039': '出生日期',
                'axcp0004': '证件类型',
                'axcp0005': '证件号码',
                'albp0289': '服务时长',
                'axcp0077': '注册状态'
            })
            print(f"数据列名: {list(detail_data.columns)}")
        except Exception as e:
            print(f"重命名列时出错: {e}")
            traceback.print_exc()
            raise

        # 确保服务时长列是数值类型
        try:
            print("转换服务时长列为数值类型")
            detail_data['服务时长'] = dd.to_numeric(detail_data['服务时长'], errors='coerce').fillna(0.0)
        except Exception as e:
            print(f"转换服务时长类型时出错: {e}")
            traceback.print_exc()
            detail_data['服务时长'] = dd.from_pandas(pd.Series([0.0] * len(detail_data)), npartitions=1)

        # 计算年龄 - 使用向量化函数替代递归
        try:
            print("开始计算年龄...")
            detail_data['年龄'] = detail_data['出生日期'].map_partitions(
    vectorized_calculate_age, meta=('年龄', 'float64')
)

            # 统计年龄计算结果
            total_records = len(detail_data)
            valid_ages = (~detail_data['年龄'].isnull()).sum().compute()
            print(f"总记录数: {total_records}, 成功计算年龄: {valid_ages}")

            # 年龄分布统计
            age_stats = detail_data['年龄'].describe().compute()
            print(f"年龄分布统计:\n{age_stats}")

            # 年龄分组
            detail_data['年龄段'] = detail_data['年龄'].map_partitions(
                lambda x: x.apply(categorize_age), meta=('年龄段', 'object')
            )

        except Exception as e:
            print(f"计算年龄时出错: {e}")
            traceback.print_exc()
            detail_data['年龄'] = dd.from_pandas(pd.Series([None] * len(detail_data)), npartitions=1)
            detail_data['年龄段'] = dd.from_pandas(pd.Series(['未知年龄'] * len(detail_data)), npartitions=1)

        # 标记证件类型相关字段
        try:
            print("标记证件类型相关字段...")
            detail_data['身份证类型'] = detail_data['证件类型'].map_partitions(
                lambda x: x.apply(lambda y: '身份证' if y == '01' else '非身份证'),
                meta=('身份证类型', 'object')
            )

            detail_data['证件号缺失'] = detail_data['证件号码'].isna() | (detail_data['证件号码'] == '')
            detail_data['有效身份证件'] = (detail_data['证件类型'] == '01') & (detail_data['证件号码'].str.len() == 18)
            detail_data['非身份证件'] = detail_data['证件类型'] != '01'
            detail_data['服务时长为0'] = (detail_data['服务时长'] == 0.0) | detail_data['服务时长'].isna()

        except Exception as e:
            print(f"标记证件类型时出错: {e}")
            traceback.print_exc()
            detail_data['身份证类型'] = dd.from_pandas(pd.Series(['未知'] * len(detail_data)), npartitions=1)
            detail_data['证件号缺失'] = dd.from_pandas(pd.Series([False] * len(detail_data)), npartitions=1)
            detail_data['有效身份证件'] = dd.from_pandas(pd.Series([False] * len(detail_data)), npartitions=1)
            detail_data['非身份证件'] = dd.from_pandas(pd.Series([False] * len(detail_data)), npartitions=1)
            detail_data['服务时长为0'] = dd.from_pandas(pd.Series([True] * len(detail_data)), npartitions=1)

        # 读取服务记录表kdx07
        try:
            print("开始读取服务记录数据...")
            service_columns = ['albp0029', 'axcp0005', 'albx0702', 'albx0704']
            service_query = f"SELECT {', '.join(service_columns)} FROM kdx07"
            
            service_dfs = []
            for i, chunk in enumerate(pd.read_sql(service_query, engine, chunksize=chunksize)):
                print(f"已读取服务记录数据块 {i+1}")
                service_dfs.append(dd.from_pandas(chunk, npartitions=1))
            
            if not service_dfs:
                print("警告: 无法从kdx07表中读取任何数据，服务时长统计将不准确")
                service_data = None
            else:
                service_data = dd.concat(service_dfs)
                # 重命名列以便于连接
                service_data = service_data.rename(columns={
                    'albp0029': '用户ID',
                    'axcp0005': '证件号码',
                    'albx0702': '服务时长',
                    'albx0704': '服务日期'
                })
                
                # 确保服务时长是数值类型
                service_data['服务时长'] = dd.to_numeric(service_data['服务时长'], errors='coerce')
                
                # 提取服务年份
                service_data['服务年份'] = service_data['服务日期'].map_partitions(
                    lambda x: x.apply(extract_year), meta=('服务年份', 'float64')
                )
                
                # 过滤掉服务时长为0或空的记录
                service_data = service_data[service_data['服务时长'] > 0]
                
                print(f"成功读取{len(service_data)}条有效服务记录")
        except Exception as e:
            print(f"读取服务记录数据时出错: {e}")
            traceback.print_exc()
            service_data = None

        # 按年龄段分组统计
        try:
            print("开始按年龄段分组统计...")
            
            # 计算重复证件号码
            duplicate_ids = detail_data.groupby('证件号码').size().reset_index()
            duplicate_ids = duplicate_ids.rename(columns={0: 'count'})
            duplicate_ids = duplicate_ids[duplicate_ids['count'] > 1]

            # 标记重复证件
            detail_data['is_duplicate'] = detail_data['证件号码'].isin(duplicate_ids['证件号码'])

            # 分组统计基础数据
            age_group_stats = detail_data[detail_data['注册状态'] == '0'].groupby('年龄段').agg({
                '年龄段': 'count',
                'is_duplicate': 'sum',
                '证件号缺失': 'sum',
                '有效身份证件': 'sum',
                '非身份证件': 'sum',
                '服务时长为0': 'sum'
            }).compute()

            # 重命名列
            age_group_stats = age_group_stats.rename(columns={
                '年龄段': '年龄段计数',
                'is_duplicate': '证件号码重复数',
                '证件号缺失': '证件号缺失数',
                '有效身份证件': '有效身份证件用户数',
                '非身份证件': '非身份证件用户数',
                '服务时长为0': '服务时长为0的用户数'
            }).reset_index()

            # 如果有服务记录数据，计算2025年前后的服务时长大于0的用户数
            if service_data is not None:
                print("开始计算按年龄段的服务时长统计...")
                
                # 合并用户数据和服务数据
                merged_data = dd.merge(
                    detail_data[['证件号码', '年龄段', '注册状态']],
                    service_data,
                    on='证件号码',
                    how='inner'
                )
                
                # 过滤注册状态为0的用户
                valid_users = merged_data[merged_data['注册状态'] == '0']
                
                # 按年龄段和服务年份统计用户数
                after_2025_users = valid_users[valid_users['服务年份'] > 2025]
                before_2025_users = valid_users[valid_users['服务年份'] < 2025]
                
                # 按年龄段分组统计
                after_2025_stats = after_2025_users.groupby('年龄段')['证件号码'].nunique().compute().reset_index()
                before_2025_stats = before_2025_users.groupby('年龄段')['证件号码'].nunique().compute().reset_index()
                
                after_2025_stats = after_2025_stats.rename(columns={'证件号码': '2025年之后服务时长大于0用户数'})
                before_2025_stats = before_2025_stats.rename(columns={'证件号码': '2025年之前服务时长大于0用户数'})
                
                # 合并到年龄分组统计结果中
                age_group_stats = pd.merge(
                    age_group_stats, 
                    after_2025_stats, 
                    on='年龄段', 
                    how='left'
                ).fillna(0)
                
                age_group_stats = pd.merge(
                    age_group_stats, 
                    before_2025_stats, 
                    on='年龄段', 
                    how='left'
                ).fillna(0)
                
                # 转换为整数类型
                age_group_stats['2025年之后服务时长大于0用户数'] = age_group_stats['2025年之后服务时长大于0用户数'].astype(int)
                age_group_stats['2025年之前服务时长大于0用户数'] = age_group_stats['2025年之前服务时长大于0用户数'].astype(int)
                
                print("按年龄段的服务时长统计完成")
            else:
                # 如果没有服务记录数据，填充默认值
                age_group_stats['2025年之后服务时长大于0用户数'] = 0
                age_group_stats['2025年之前服务时长大于0用户数'] = 0
                print("使用默认值填充服务时长统计（无服务记录数据）")

        except Exception as e:
            print(f"分组统计时出错: {e}")
            traceback.print_exc()
            age_group_stats = pd.DataFrame()

        # 输出统计结果
        print("\n基本统计结果:")
        for name, value in basic_stats.items():
            if value is not None:
                print(f"{name}: {value}")
            else:
                print(f"{name}: 统计失败")

        print("\n年龄分组统计结果:")
        if not age_group_stats.empty:
            print(age_group_stats)
            print(f"年龄分组统计总人数: {age_group_stats['年龄段计数'].sum()}")
        else:
            print("年龄分组统计失败，结果为空")

        # 保存结果到Excel
        if output_file and not age_group_stats.empty:
            try:
                print(f"保存结果到Excel: {output_file}")
                with pd.ExcelWriter(output_file) as writer:
                    pd.DataFrame(basic_stats.items(), columns=['统计项', '数量']).to_excel(
                        writer, sheet_name='基本统计', index=False)
                    age_group_stats.to_excel(writer, sheet_name='年龄分组统计', index=False)
                print(f"\n结果已成功保存到: {output_file}")
            except Exception as e:
                print(f"保存结果到Excel时出错: {e}")
                traceback.print_exc()

        # 关闭数据库连接
        if connection:
            connection.close()

        return {
            'basic_stats': basic_stats,
            'age_group_stats': age_group_stats
        }

    except Exception as e:
        print(f"数据分析过程中发生错误: {e}")
        traceback.print_exc()
        return None



if __name__ == "__main__":
    # 数据库配置 - PostgreSQL
    database_config = {
        'type': 'postgresql',
        'host': '192.168.29.130',
        'user': 'postgres',
        'password': '123456',
        'database': 'manuals',
        'port': 5432,
        'schema': 'public'
    }

    output_file = r'C:\Users\17610\Desktop\数据清洗\数据统计\统计结果.xlsx'

    print("开始分析数据...")
    results = analyze_data(database_config, output_file)
    if results:
        print("数据分析完成!")